# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import enum
import functools
import typing

import requests
import tornado.gen
import cachetools

import coin.flow.topic_map as topic_map
import coin.flow.subscriber as flow_subscriber
import coin.proto.coin_market_enums_pb2 as coin_enum
import coin.proto.coin_query_pb2 as coin_query

from coin.base.config import Config
from coin.base.timestamp import get_timestamp
from coin.exchange.base.order_fill_util import FillManager
from coin.exchange.base.rest.task_scheduler import Scheduler, wrap_async_task
from coin.exchange.base.rest_order_gateway_base import PrivateExchangeMessage
from coin.exchange.bitstamp_v2.kr_rest.private_client import BitstampPrivateParser
from coin.exchange.util.rest_rate_limit import RestRateLimiter
from coin.proto.coin_order_gateway_pb2 import (
    OrderGatewayConfig,
    OrderEvent,
)
from coin.exchange.bitstamp_v2.kr_rest.product import (
    BitstampProduct,
    BitstampCurrency,
)
from coin.exchange.bitstamp_v2.kr_rest.native_private_client import (
    BitstampNativePrivateClient,
    TransactionType,
    BitstampOrderSide,
)
from coin.exchange.bitstamp_v2.ws.private_subscriber import (
    BitstampWsPrivateClient,
    ChannelType,
    EventType,
    parse_channel_name,
)
from coin.exchange.base.rest_order_gateway_base import (
    Order,
    OrderType,
    OrderSide,
    BaseBalanceOrderGatewayInfo,
    RestOrderGateway,
    create_order_internal,
    exchange_message_handler,
    to_proto_order_side,
)


def create_order_from_open_order_message(message, timestamp, order_id=None):
  """
  rest:
  {
    "price": "169.54",
    "currency_pair": "ETH/USD",
    "datetime": "2019-04-24 03:25:31",
    "amount": "0.10000000",
    "type": "1",
    "id": "3174363621"
  }
  ws:
  {
    "data": {
      "microtimestamp": "1555919888674561",
      "amount": 0.01,
      "order_type": 1,
      "amount_str": "0.01000000",
      "price_str": "5808.71",
      "price": 5808.71,
      "id": 3166469679,
      "datetime": "1555919888"
    },
    "event": "order_created",
    "channel": "live_orders_btcusd"
  }
  """
  product = BitstampProduct.FromStrNativeProduct(message['currency_pair'])
  price_orig = float(message['price'])
  qty_orig = float(message['amount'])
  order_side = getattr(OrderSide, BitstampOrderSide(int(message['type'])).name)
  exchange_order_id = str(message['id'])
  status = message.get('status')
  order_type = OrderType.LIMIT

  fill = FillManager(
      price_orig=price_orig,
      qty_orig=qty_orig,
      price_last_fill=None,
      qty_last_fill=None,
      price_avg_fill=None,
      qty_accum_fill=0,
      timestamp_last_fill=None,
  )
  order = Order(order_id=order_id,
                product=product,
                price=price_orig,
                qty=qty_orig,
                order_side=order_side,
                order_type=order_type,
                tag=None,
                accepted=True,
                cancel_sent=False,
                internal=create_order_internal(
                    timestamp=timestamp,
                    exchange_order_id=exchange_order_id,
                    fill=fill,
                    status=status,
                ))
  return order


class BitstampOrderGatewayInfo(BaseBalanceOrderGatewayInfo):
  CurrencyType = BitstampCurrency
  ProductType = BitstampProduct
  BalanceType = float
  PositionType = float

  def update_balance(self, message):
    msg = BitstampPrivateParser.parse_account_balance(message)
    for balance in msg.each_balance:
      currency = BitstampCurrency.FromStr(balance.currency)
      if currency in self.currencies:
        self.set_balance(currency, balance.total, balance)


class BitstampOrderGateway(RestOrderGateway):
  _exchange = 'Bitstamp'
  _market_type = 'Spot'
  _api_version = None

  OGInfoType = BitstampOrderGatewayInfo

  class _MsgType(enum.IntEnum):
    REST_ACCOUNT_BALANCE = 1
    REST_TRANSACTIONS = 2
    REST_OPEN_ORDERS = 3
    WS_LIVE_ORDERS = 4
    WS_FILLS = 5

  def __init__(self,
               currencies: typing.List[BitstampCurrency],
               products: typing.List[BitstampProduct],
               config: Config,
               og_config=None,
               logger=None):
    super().__init__(currencies, products, config, og_config, logger)

    # Following instance attributes are initialized in start method.
    self._flow_sub = None
    self._ioloop = None
    self._publisher = None
    self._scheduler = None
    self._private_client = None
    self._ws_private_client = None
    self._rest_rate_limiter = RestRateLimiter(delay=10)
    self._fill_ttl_cache = cachetools.TTLCache(maxsize=20, ttl=10)

    self._og_config = OrderGatewayConfig(
        account_query_period=1,
        orders_snapshot_query_period=5,
        fill_query_period=2,
    )
    if og_config is not None:
      self._og_config.MergeFrom(og_config)
    self._og_logger.exchange_type = coin_enum.Bitstamp

  def submittable(self):
    return self.is_ready() and self._rest_rate_limiter.is_available()

  def start(self, queue, ioloop, flow_sub=None):
    api_host = self._og_config.rest_root_url if \
               self._og_config.HasField('rest_root_url') else None
    disable_ssl = self._og_config.rest_disable_ssl

    topic_string = 'bitstamp_rest_private_1'
    topic_id = topic_map.search_topic_id(topic_string)[0][1]
    self._ioloop = ioloop

    self._publisher = queue.get_writer(topic_id=topic_id)
    self._scheduler = Scheduler(ioloop=ioloop, logger=self._logger)

    self._private_client = BitstampNativePrivateClient(
        key_file=self._config.key_file,
        api_host=api_host,
        disable_ssl=disable_ssl,
        use_async=True,
        max_connections=self._og_config.max_connections,
        rest_request_gap_sec=self._og_config.rest_request_gap_sec,
    )
    self._prepare_ws_private_client(start=True)
    self._install_periodic_query()

    # Subscribe to flow for feed.
    self._flow_sub = flow_sub or flow_subscriber.from_queue(queue)
    self._flow_sub.subscribe(topic_string, self.on_private_exchange_msg)

  @tornado.gen.coroutine
  def _install_periodic_query(self):
    assert self._scheduler is not None, 'self._scheduler is not set!'
    assert self._publisher is not None, 'self._publisher is not set!'

    # query balance
    query_balance_task = wrap_async_task(
        self._private_client.query_account_balance,
        msg_type=self._MsgType.REST_ACCOUNT_BALANCE,
        publisher=self._publisher,
    )
    self._scheduler.add_periodic_task(
        period=self._og_config.account_query_period,
        job_func=query_balance_task,
    )

    yield tornado.gen.sleep(self._og_config.burst_query_avoid_time)
    # query open orders
    query_orders_task = wrap_async_task(
        self._private_client.query_open_orders,
        msg_type=self._MsgType.REST_OPEN_ORDERS,
        publisher=self._publisher,
    )
    self._scheduler.add_periodic_task(period=self._og_config.orders_snapshot_query_period,
                                      job_func=query_orders_task)

    yield tornado.gen.sleep(self._og_config.burst_query_avoid_time)
    # query fills
    query_orders_info_task = wrap_async_task(
        self._private_client.query_user_transactions,
        msg_type=self._MsgType.REST_TRANSACTIONS,
        publisher=self._publisher,
    )
    self._scheduler.add_periodic_task(
        period=self._og_config.fill_query_period,
        job_func=query_orders_info_task,
    )
    self.start_scheduler()

  @tornado.gen.coroutine
  def _prepare_ws_private_client(self, start=True):
    if self._ws_private_client is not None:
      self._ws_private_client.stop()

    yield tornado.gen.sleep(5)
    close_callback = functools.partial(self._prepare_ws_private_client, True)
    self._ws_private_client = BitstampWsPrivateClient(
        ioloop=self._ioloop,
        notification_callback=self._notification_callback,
        close_callback=close_callback,
        on_ready=self._on_ws_private_client_ready,
    )
    if start:
      self._ws_private_client.start()

  def _notification_callback(self, ws_msg):
    timestamp = get_timestamp()
    channel = parse_channel_name(ws_msg['channel'])['channel_type']
    if channel == ChannelType.ORDERS:
      msg_type = self._MsgType.WS_LIVE_ORDERS
    elif channel == ChannelType.TRADE:
      msg_type = self._MsgType.WS_FILLS
    else:
      self._logger.error('Unexpected ws message: %s', ws_msg)
      return

    exchange_message = PrivateExchangeMessage(msg_type=msg_type,
                                              data=ws_msg,
                                              received_timestamp=timestamp,
                                              request_params=None,
                                              request_timestamp=None)
    self._publisher.write(timestamp, exchange_message)

  @tornado.gen.coroutine
  def _on_ws_private_client_ready(self):
    channels = [ChannelType.TRADE, ChannelType.ORDERS]
    for product in self._og_info.products:
      for channel in channels:
        yield self._ws_private_client.subscribe_channel(channel, product.native_symbol)
        self._logger.info('Subscribe to %s, %s', channel, product)

  def _async_submit_order(self, order, post_only=None):
    if post_only is True:
      self._logger.debug('Post-only order is not supported!')

    formatter = self._order_formatter[order.product]
    params = {
        'currency_pair': order.product.native_symbol,
        'price': formatter.format_price(order.price),
        'amount': formatter.format_qty(order.qty),
    }
    if order.order_side == OrderSide.BUY:
      fut = self._private_client.buy_limit_order(**params)
    elif order.order_side == OrderSide.SELL:
      fut = self._private_client.sell_limit_order(**params)
    else:
      raise ValueError('Order side %s not supported!', order.order_side)
    return fut

  def _async_cancel_order(self, order):
    assert order.internal.exchange_order_id is not None, order
    fut = self._private_client.cancel_order(order_id=order.internal.exchange_order_id)
    return fut

  def _handle_submit_response(self, order, response):
    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      self._order_manager.remove_order(order, ignore_error=True)
      return OrderEvent.ORDER_REJECTED

    message = response.json()
    status = message.get('status')
    if status is not None:  # == "error"
      self._logger.error('Order submission failed! %s', message)
      self._order_manager.remove_order(order, ignore_error=True)
      return OrderEvent.ORDER_ERROR

    exchange_order_id = message.get('id')
    if exchange_order_id is None:
      self._logger.error('Weird! Accepted order without order id! %s', message)
      self._order_manager.remove_order(order, ignore_error=True)
      return OrderEvent.ORDER_REJECTED

    prev_order = self._order_manager.get(exchange_order_id=exchange_order_id)
    if prev_order is None:
      self._order_manager.set_order(order, exchange_order_id=exchange_order_id)
    else:
      if prev_order.order_id != order.order_id:
        self._logger.info('order id replaced: %s => %s', prev_order.order_id, order.order_id)
        self._order_manager.remove_order(prev_order)
        self._order_manager.set_order(order, exchange_order_id=exchange_order_id)
        self._order_manager.remove_order(order)
        prev_order.order_id = order.order_id
        self._order_manager.add_order(prev_order)
    return OrderEvent.ORDER_ACCEPTED

  def _handle_cancel_response(self, order, response):
    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      if response.status_code == 403:
        message = response.json()
        if message.get('code') == 'API0004':  # Invalid nonce
          order.cancel_sent = False
      return OrderEvent.CANCEL_ERROR

    message = response.json()
    error = message.get('error')
    if error is not None:
      if error == 'Order in queue. Please try again later.':
        order.cancel_sent = False
      else:
        # (NotFound case)
        pass
      self._logger.error('Cancel failed! %s', message)
      return OrderEvent.CANCEL_ERROR

    exchange_order_id = str(message['id'])
    if exchange_order_id is None:
      self._logger.error('Weird! Cancel accepted order without order id! %s', message)
      order.cancel_sent = False
      return OrderEvent.CANCEL_ERROR
    elif exchange_order_id != order.internal.exchange_order_id:
      order.cancel_sent = False
      self._logger.error('Insane response, exchange order id not consistent: %s, %s',
                         order,
                         message)
      return OrderEvent.CANCEL_ERROR
    else:
      if not self._order_manager.is_dead_order(order):
        self._order_manager.set_order(order, cancel_confirmed=True)
      else:
        self._logger.info('order is canceled already! %s', order.order_id)
      return OrderEvent.CANCEL_CONFIRMED

  def _handle_error_response(self, response):
    if response.status_code in (503, 504):
      self._logger.error('503/504 Gateway Timeout error!')
    elif response.status_code == 429:
      self._rest_rate_limiter.set_error(response.status_code)
      self._logger.error('Rate limit hit (429)!')

    self._logger.error('response error! %s, %s, %s',
                       response.status_code,
                       response.content,
                       response.url)

  def _handle_single_fill(self, fill_info, timestamp):
    # noinspection SpellCheckingInspection
    fill_id = str(fill_info['id'])
    if fill_id in self._processed_fills:
      return

    transaction_type = TransactionType(int(fill_info['type']))
    if transaction_type != TransactionType.MARKET_TRADE:
      self._logger.info('Balance change due to non-trade events! %s, %s',
                        fill_info,
                        self._processed_fills)
      self._processed_fills.add(fill_id)
      return

    exchange_order_id = str(fill_info['order_id'])
    order = self._order_manager.get(exchange_order_id=exchange_order_id)
    if order is None:
      self._logger.info('Untrackable fill received! %s', fill_info)
      return

    native_base_currency = order.product.base.native_currency
    native_quote_currency = order.product.quote.native_currency
    pair = '%s_%s' % (native_base_currency, native_quote_currency)
    base_qty = abs(float(fill_info[native_base_currency]))
    price_last_fill = float(fill_info[pair])
    fill_type = self._guess_fill_type(order, price_last_fill, timestamp)
    last_fill = order.internal.fill.update_by_last_fill(price_last_fill, base_qty)
    if last_fill:
      fill_proto = coin_query.ProductFillElement(fill_id=fill_id,
                                                 price=last_fill.price_last_fill,
                                                 qty=last_fill.qty_last_fill,
                                                 order_id=order.order_id,
                                                 side=to_proto_order_side(order.order_side),
                                                 exchange_order_id=exchange_order_id,
                                                 guess_fill_type=fill_type,
                                                 symbol=order.product.symbol)
      self.process_fill(fill_proto, timestamp)
    else:
      self._logger.error('Filled qty more than the orig_qty: %s', order.order_id)
    self._processed_fills.add(fill_id)

  @exchange_message_handler(_MsgType.REST_OPEN_ORDERS)
  def _handle_rest_open_orders(self, exchange_message: PrivateExchangeMessage):
    timestamp = get_timestamp()
    response = exchange_message.data
    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      return

    message = response.json()
    exchange_open_order_dict = {}
    for entry in message:
      native_symbol = entry['currency_pair']
      product = BitstampProduct.FromStr(native_symbol)
      if product not in self._og_info.products:
        continue  # ignore

      order = create_order_from_open_order_message(entry, timestamp)
      exchange_order_id = order.internal.exchange_order_id
      prev_order = self._order_manager.get(exchange_order_id=exchange_order_id)
      if prev_order is None:
        order_id = self._order_id_generator.gen_order_id(is_foreign=True)
        order.order_id = order_id
      else:
        order.order_id = prev_order.order_id
      exchange_open_order_dict[exchange_order_id] = order
    self._order_manager.copy_from_exchange_open_orders(exchange_order_dict=exchange_open_order_dict,
                                                       products=self._og_info.products)

    if not self.is_ready():
      self._og_ready.set_orders_ready(True)
      self._og_ready.check_ready()

  @exchange_message_handler(_MsgType.REST_ACCOUNT_BALANCE)
  def _handle_rest_balance(self, exchange_message: PrivateExchangeMessage):
    response = exchange_message.data
    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      return

    message = response.json()
    self._og_info.update_balance(message)
    self._og_logger.write_balance(self._og_info.balance_as_proto())
    if not self.is_ready():
      self._og_ready.set_balance_ready(True)
      self._og_ready.check_ready()

  @exchange_message_handler(_MsgType.REST_TRANSACTIONS)
  def _handle_rest_fills(self, exchange_message: PrivateExchangeMessage):
    response = exchange_message.data
    timestamp = exchange_message.received_timestamp
    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      return

    message = response.json()
    is_first_query = not bool(self._processed_fills)
    if is_first_query:
      self._processed_fills.add(None)
      for fill_info_msg in message:
        fill_id = str(fill_info_msg['id'])
        self._processed_fills.add(fill_id)
      return

    for fill_info_msg in message:
      self._handle_single_fill(fill_info_msg, timestamp)

  @exchange_message_handler(_MsgType.WS_LIVE_ORDERS)
  def _handle_ws_orders(self, exchange_message: PrivateExchangeMessage):
    timestamp = exchange_message.received_timestamp
    data = exchange_message.data
    native_symbol = parse_channel_name(data['channel'])['currency_pair']
    data['data']['currency_pair'] = native_symbol
    data['data']['status'] = EventType(data['event'])
    data['data']['type'] = data['data']['order_type']
    order_info = create_order_from_open_order_message(data['data'], timestamp)
    exchange_order_id = order_info.internal.exchange_order_id
    order = self._order_manager.get(exchange_order_id=exchange_order_id)
    if order is None:
      return

    if order_info.internal.status == EventType.ORDER_DELETED:
      if not self._order_manager.is_dead_order(order):
        self._logger.info('Treat as canceled order! %s', order.order_id)
        self._order_manager.set_order(order, cancel_confirmed=True)

  @exchange_message_handler(_MsgType.WS_FILLS)
  def _handle_ws_fills(self, exchange_message: PrivateExchangeMessage):
    data = exchange_message.data
    fill_id = str(data['data']['id'])
    # Assume there is no self trade.
    buy_order_id = str(data['data']['buy_order_id'])
    order = self._order_manager.get(exchange_order_id=buy_order_id)
    if order is None:
      sell_order_id = str(data['data']['sell_order_id'])
      order = self._order_manager.get(exchange_order_id=sell_order_id)
      if order is None:
        self._fill_ttl_cache[fill_id] = exchange_message
      return

    native_symbol = parse_channel_name(data['channel'])['currency_pair']
    product = BitstampProduct.FromStrNativeProductNothrow(native_symbol)

    if product != order.product:
      self._logger.error('product does not match. %s, %s', order, data)
      return
